import os

import pandas as pd
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import JsonOutputKeyToolsParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_experimental.tools import PythonAstREPLTool
from langchain_ollama import ChatOllama


def code_print(res):
    print("即将运行python代码", res)
    return res

def safe_extract_query(parsed_result):
    print("🧪 parser 输出：", parsed_result)
    if not isinstance(parsed_result, list) or len(parsed_result) == 0:
        raise ValueError("❌ parser 输出不是非空列表，LLM 可能没有按要求返回 query 结构。")
    item = parsed_result[0]
    if "query" not in item:
        raise ValueError(f"❌ query 字段不存在。实际内容: {item}")
    return item["query"]

if __name__ == '__main__':
    load_dotenv(override=True)

    df = pd.read_csv("archive/WA_Fn-UseC_-Telco-Customer-Churn.csv")
    tool = PythonAstREPLTool(locals={"df": df})
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    # model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    model = ChatOllama(model="qwen3:latest", base_url="http://192.168.97.217:11434")
    llm_with_tools = model.bind_tools([tool])

    # response = llm_with_tools.invoke("我有一张表名为'df',请帮我计算MonthlyCharges字段的均值。")
    parser = JsonOutputKeyToolsParser(key_name=tool.name, firt_tool_only=True)
    # llm_chain = llm_with_tools | parser
    # response = llm_chain.invoke("我有一张表名为'df',请帮我计算MonthlyCharges字段的均值。")
    # print(response)
    system = f"""你可以访问一个名为'df'的pandas数据框，你可以使用df.head().to_markdown() 查看数据集的基本信息，请根据用户提出的问题，
    编写python代码来回答。只返回代码，不返回其他内容。只允许使用pandas和内置库。"""
    prompt = ChatPromptTemplate(
        [("system", system),
         ("user", "{question}")]
    )
    # llm_chain = prompt | llm_with_tools | parser
    # response = llm_chain.invoke("我有一张表名为'df',请帮我计算MonthlyCharges字段的均值。")
    # print(response)
    extract_query = RunnableLambda(safe_extract_query)
    code_print = RunnableLambda(code_print)
    code_chain = prompt | llm_with_tools | parser | extract_query | code_print | tool
    # response = code_chain.invoke({"question": "请帮我计算MonthlyCharges字段的均值。"})
    # print(response)
    response = code_chain.invoke({"question": "请帮我分析gender和SeniorCitizen和Churn三个字段之间的关系"})
    # response = code_chain.invoke({"questin": "请帮我计算MonthlyCharges字段的均值。"})
    print(response)
